import open3d as o3d
import copy
import numpy as np

def shift_point_cloud(pc, shift_range=0.1):
    N, C = pc.points.shape
    shifts = np.random.uniform(-shift_range, shift_range, (1, C)).astype(np.float32)
    pc.points += shifts
    return pc

def draw_registration_result(source, target, transformation):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.paint_uniform_color([1, 0.706, 0])
    target_temp.paint_uniform_color([0, 0.651, 0.929])
    source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp],
                                      zoom=0.4459,
                                      front=[0.9288, -0.2951, -0.2242],
                                      lookat=[1.6784, 2.0612, 1.4451],
                                      up=[-0.3402, -0.9189, -0.1996])
    

if __name__=='__main__':
    import cv2
    from matplotlib import pyplot as plt
    
    im=cv2.imread('/home/wsl/gitee/3D/data/dianchi/im/92801 (2).npy.png')
    plt.imshow(im)
    plt.show()
    # source = o3d.io.read_point_cloud("/home/wsl/zz/replay_93936_2021-9-22.ply")
    # target = o3d.io.read_point_cloud("/home/wsl/gitee/3D/data/temp/replay_99587_2021-9-24.npy")
    # #target=shift_point_cloud(target)
    # threshold = 1
    # # trans_init = np.asarray([[0.862, 0.011, -0.507, 0.5],
    # #                          [-0.139, 0.967, -0.215, 0.7],
    # #                          [0.487, 0.255, 0.835, -1.4],
    # #                          [0.0, 0.0, 0.0, 1.0]])
    # trans_init = np.asarray([[1, 0.0, 0.0, 0.0],
    #                          [0.0, 1.0, 0.0, 0.0],
    #                          [0.0, 0.0, 1.0, 0.0],
    #                          [0.0, 0.0, 0.0, 1.0]])
    # #draw_registration_result(source, target, trans_init)
    #
    # # print("Apply point-to-point ICP")
    # # reg_p2p = o3d.pipelines.registration.registration_icp(
    # #     source, target, threshold, trans_init,
    # #     o3d.pipelines.registration.TransformationEstimationPointToPoint())
    # # print(reg_p2p)
    # # print("Transformation is:")
    # # print(reg_p2p.transformation)
    # # draw_registration_result(source, target, reg_p2p.transformation)
    #
    # reg_p2p = o3d.pipelines.registration.registration_icp(source, target, threshold, trans_init,
    #                                             o3d.pipelines.registration.TransformationEstimationPointToPoint(),
    #                                             o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=30))
    # #reg_p2p = o3d.pipelines.registration.registration_ransac_based_on_correspondence(source,target,)
    # print(reg_p2p)
    # print("Transformation is:")
    # print(reg_p2p.transformation)
    # draw_registration_result(source, target, reg_p2p.transformation)